# Description: Here is Main function os LinearDesign Algorithm implementated with Python
# Author: Shibo Li, MiQroEra
# Date: 2023-07-15
# Version: 1.0
# I am not sure if the optimized mRNA Sequence fits the reality or not.
# It requires further discussion.


import argparse
import json
import pandas as pd
from utils.codon_lattice import generate_codon_lattice
from utils.optimize_sequence import find_optimal_path
from utils.evaluate import calculate_cai, calculate_mfe
from utils.welcome import welcome_page

parser = argparse.ArgumentParser(description='mRNA Optimization')
parser.add_argument('--codon_table', type=str, default='./data/codon_table.json',
                    help='Path to the codon table JSON file (default: ./data/codon_table.json)')
parser.add_argument('--codon_usage', type=str, default='./data/codon_usage.csv',
                    help='Path to the codon usage CSV file (default: ./data/codon_usage.csv)')
parser.add_argument('--protein_sequence', type=str, default='MADLWQLLLTLALAGSSDAFSGSEATAAILSRAPWSLQSVNP',
                    help='Target protein sequence (default: MADLWQLLLTLALAGSSDAFSGSEATAAILSRAPWSLQSVNP)')
args = parser.parse_args()

with open(args.codon_table) as f:
    codon_table = json.load(f)['codon_table']

codon_usage_df = pd.read_csv(args.codon_usage)
codon_usage = codon_usage_df.set_index('Codon')['Probability'].to_dict()

welcome_page()

print("如果您没有使用--protein_sequence参数来输入其他序列，我们将会使用默认序列来展示程序功能。")
print("我们的默认序列为：MADLWQLLLTLALAGSSDAFSGSEATAAILSRAPWSLQSVNP")
print("您可以在--protein_sequence参数后输入您想要的序列。")
print("您可以在--codon_table参数后输入您想要的密码子表。")

print("*" * 50)

codon_lattice = generate_codon_lattice(codon_table, args.protein_sequence)

optimized_mrna_sequence = find_optimal_path(codon_lattice, codon_usage)
print("您输入的protein_sequence为：", args.protein_sequence)
print("优化后的mRNA序列：", optimized_mrna_sequence)


optimized_cai = calculate_cai(optimized_mrna_sequence, codon_usage)
optimized_mfe = calculate_mfe(optimized_mrna_sequence)

print("优化后的mRNA序列的CAI：", optimized_cai)
print("优化后的mRNA序列的MFE：", optimized_mfe, "kcal/mol/nt")
